RL: Training Inference Mismatch

导言

• 25年，RL训练崩溃归因于训推不一致；
• 为此提出了很多方法，TIS，Router Replay，FP16训推，batch一致性…
• 如何判断 模型当前训推不一致，并找到不一致实现处，是实践的要点。

基本概念

KL 散度

Reverse KL Divergence

GRPO 使用的是 Reverse KL（采样自当前策略 π_θ，参考策略为 π_ref）：

$$D_{KL}(\pi_\theta | \pi_{ref}) = \mathbb{E}{x \sim \pi_\theta}\left[\log\frac{\pi_\theta(x)}{\pi{ref}(x)}\right]$$

在 LLM 场景下，按 token 级别计算：

KL = Σ_t π_θ(o_t | q, o_<t) · log[π_θ(o_t | q, o_<t) / π_ref(o_t | q, o_<t)]

三种 KL 估计方法

由于直接计算 KL 的期望成本高，实践中采用蒙特卡洛近似。常见有三种估计器 [[49]][[50]]：

估计器	公式	特性	适用场景
k1	-log(r)，其中 r = π_ref/π_θ	无偏但方差极大，梯度不含 π_ref	PPO 中作为 reward shaping，不适合作为独立 KL loss
k2	0.5 * (log(r))²	有偏但方差低，梯度等价于 Reverse KL	✅ GRPO 推荐（VeRL/TRL 默认）
k3	(r - 1) - log(r)	无偏、方差低，但梯度等价于 Forward KL	需注意：采样分布不匹配时可能不稳定

关键代码逻辑（VeRL/TRL 实现）：

# 假设已获取 per-token logps
log_ratio = ref_logps - actor_logps  # log(π_ref/π_θ)

# k1: 原始 KL (高方差)
kl_k1 = -log_ratio

# k2: 平方近似 (低方差，推荐)
kl_k2 = 0.5 * (log_ratio ** 2)

# k3: Bregman 形式 (无偏，但梯度对应 Forward KL)
kl_k3 = (log_ratio.exp() - 1) - log_ratio

# 最终 loss 加入
loss = policy_loss + beta * kl_loss_type(actor_logps, ref_logps)

🔍 为什么 k2 更推荐？

• k2 的梯度：∇_θ [0.5*(log r)²] = (log r) · ∇_θ log π_θ，恰好匹配 Reverse KL 的实用梯度形式 [[49]]
• k3 虽然无偏，但其梯度对应 Forward KL，在 π_θ 与 π_ref 差距较大时，重要性采样权重 π_ref/π_θ 可能爆炸，导致训练不稳定

监控指标的计算流程（每步 RL）

1️⃣ 采样阶段：
   - 对每个 query，用当前策略 π_θ 生成 G 个 completions
   - 同时用参考策略 π_ref 计算相同 tokens 的 log-probs

2️⃣ 计算 per-token KL：
   log_ratio = log(π_ref) - log(π_θ)
   kl_token = kl_loss_type(log_ratio)  # k1/k2/k3 选一

3️⃣ 聚合为标量指标：
   - 按 completion 平均：kl_per_seq = mean(kl_token over tokens)
   - 按 batch 平均：kl_monitor = mean(kl_per_seq over all completions)

4️⃣ 用于：
   - 📊 监控：TensorBoard/W&B 记录 kl 曲线，判断策略漂移
   - ⚖️ 正则：loss += β * kl_monitor（β=0.001~0.04，依任务调整）

---

四、实践建议

1. 配置选择（以 VeRL 为例）[[13]][[41]]：

   actor_rollout_ref:
     actor:
       use_kl_loss: true          # 启用 KL 正则（GRPO 必须）
       kl_loss_coef: 0.001        # β 系数，数学任务可增至 0.04
       kl_loss_type: "k2"         # 推荐 k2；若用 k3+ 可开启 straight-through 梯度修正

2. 监控阈值参考：

   • KL < 0.01：策略变化过小，可能学习缓慢
   • KL ∈ [0.01, 0.1]：健康更新区间
   • KL > 0.2：策略漂移过大，需检查 β 或奖励设计

3. 调试技巧：

   • 同时记录 kl_k2 和 kl_k3，若二者差异显著，说明 π_θ 与 π_ref 已偏离较大
   • 若 KL 持续上升且 reward 不增，考虑定期重置 reference model（DeepSeek-R1 实践）[[50]]

log p

log_p 是模型对每个位置实际生成的 token 计算出的对数概率（Log Probability），属于标量值，而 hidden_state 是 Transformer 层输出的高维稠密向量。两者在计算链路、数据形态和用途上完全不同。

---

🔍 log_p 的完整计算链路

Input Tokens 
   ↓ [Embedding]
Hidden States (layer 0) 
   ↓ [Transformer Blocks × N]
Final Hidden States: h ∈ ℝ^{B×L×D}  ← 这才是你问的 hidden_state
   ↓ [LM Head: Linear + Bias]
Logits: z ∈ ℝ^{B×L×V}            ← 词表大小 V 的未归一化得分
   ↓ [Log Softmax]
Log Probabilities: log_p_all ∈ ℝ^{B×L×V}
   ↓ [Gather 实际 token ID]
log_p ∈ ℝ^{B×L}                  ← 你问的 log_p（每个位置一个标量）

📐 维度与形态对比

概念	形状	数据类型	物理含义
hidden_state	[B, L, D]	float32/16	Transformer 输出的上下文表征向量
logits	[B, L, V]	float32/16	词表每个 token 的原始得分
log_p	[B, L]	float32/16	当前策略下，每个位置真实 token 的对数概率

💡 D 通常为 4096/7680 等，V 为词表大小（如 32k/128k），而 log_p 已坍缩到 [B, L]，只保留实际生成 token 的概率信息。

---

💻 代码级直观实现（PyTorch）

import torch
import torch.nn.functional as F

# 假设 model 已加载，input_ids 为 [B, L]
outputs = model(input_ids=input_ids, return_dict=True)
logits = outputs.logits  # [B, L, V]

# 1. 计算全词表 log softmax
log_probs_all = F.log_softmax(logits, dim=-1)  # [B, L, V]

# 2. 提取实际 token 对应的 log_p
# input_ids.unsqueeze(-1) -> [B, L, 1]
token_log_p = log_probs_all.gather(dim=-1, index=input_ids.unsqueeze(-1)).squeeze(-1)  # [B, L]

# 3. 通常只保留 response 部分（prompt 和 padding 用 mask 过滤）
response_mask = (input_ids >= tokenizer.vocab_size)  # 示例：假设 response token ID 较大
valid_log_p = token_log_p * response_mask  # 后续计算 KL/Loss 时会 mask 掉无效位置

---

⚠️ 常见误区澄清

误区	正确理解
“log_p 就是 hidden_state”	hidden_state 是向量表征；log_p 是经过 LM Head + LogSoftmax + Gather 后的标量概率
“推理时不需要 log_p”	推理（生成）时模型会隐式计算它用于采样；RL 训练时需显式保存用于 loss
“log_p 越大越好”	仅表示模型对该 token 更自信；RL 中需与 reward/advantage 结合，盲目最大化会导致 mode collapse
“KL 直接用 hidden_state 算”	KL 是概率分布距离，必须基于 log_p；hidden_state 是特征空间，无法直接算分布散度

---

📌 总结

• log_p = 每个位置真实 token 的对数概率，形状 [B, L]
• 由 hidden_state → logits → log_softmax → gather 得到，不是 hidden_state 本身
• 在 GRPO 中用于：KL 正则、importance ratio、策略梯度、loss mask
• 训练时需同时保存 actor_logp 和 ref_logp，推理时通常不显式输出但底层会计算

训推一致性

LogP Diff vs KL 散度：本质区别与阈值含义

维度	🔹 LogP Diff (训推一致性)	🔹 KL 散度 (RL 正则/监控)
比较对象	同一模型，train mode vs infer mode	两个策略，actor π_θ vs reference π_ref
核心目标	验证数值计算一致性（精度对齐）	控制策略更新幅度（防止分布崩溃）
数学形式	δ = |log_p^train - log_p^infer|	KL = 𝔼[log(π_θ/π_ref)] 或其近似
是否取绝对值	✅ 是，关注偏差大小	❌ 否，保留方向信息（谁更自信）
是否加权	❌ 所有 token 平等对待	✅ 高概率 token 贡献更大（p·log(p/q)）
量纲/单位	log 空间的相对误差（无单位比值）	信息论距离（nat，无单位但数值意义不同）
典型阈值	rel_diff < 0.01 (1%)	KL < 0.01~0.1 (依任务/β系数调整)

---

# 在 trainer 的 logging 阶段
def compute_consistency_and_kl_metrics(batch):
    metrics = {}
    
    # 🔹 训推一致性（仅调试阶段启用）
    if config.debug_consistency:
        with torch.no_grad():
            logp_train = model_train(input_ids).logps      # train mode
            logp_infer = model_infer(input_ids).logps      # eval mode
        rel_diff = (logp_train - logp_infer).abs() / (logp_train.abs() + 1e-8)
        metrics["debug/consistency_rel_diff"] = verl_F.masked_mean(
            rel_diff, batch["response_mask"]
        ).item()
    
    # 🔹 KL 散度（训练必选）
    log_ratio = batch["ref_logps"] - batch["actor_logps"]
    kl_token = 0.5 * (log_ratio ** 2)  # k2
    metrics["actor/kl_loss"] = agg_loss(
        kl_token, batch["response_mask"], config.loss_agg_mode
    ).item()
    
    return metrics

训推一致性指标

你的理解里有一个关键偏差：**RL 训练里计算 log p 时，训练端不是只输入 response，而是输入 prompt + rollout 生成的 response**。response 在这里有双重身份：

• 作为 label/action：要计算每个生成 token 的 log probability；
• 作为 teacher forcing 的上下文：第 (t) 个 response token 之后的 token，需要以前面已经生成的 response token 为条件。

所以训练和推理对齐的不是“API 传入的张量长得一样”，而是对齐同一个条件概率：

[
\log \pi_\theta(y_t \mid x, y_{<t})
]

其中：

• (x)：prompt；
• (y_t)：第 (t) 个生成出来的 response token；
• (y_{<t})：它之前已经生成的 response tokens。

推理时逐步得到这些概率；训练时用一次 causal forward 并行算出这些概率。数学上是等价的，只要 token、mask、position id、模型权重、logits 处理方式一致。

---

log p 到底对齐什么

设 prompt 为：

x = [x1, x2]

rollout 生成的 response 为：

y = [y1, y2, y3]

推理阶段并不是永远只输入 prompt。实际过程是：

prefill([x1, x2])      -> log p(y1 | x1, x2)
decode(y1)             -> log p(y2 | x1, x2, y1)
decode(y2)             -> log p(y3 | x1, x2, y1, y2)

训练阶段会把同一条轨迹拼起来：

input = [x1, x2, y1, y2, y3]

然后 causal LM 的 logits 对齐关系是：

logits at x2 -> predict y1
logits at y1 -> predict y2
logits at y2 -> predict y3

因此训练端算的是：

[
\log \pi_\theta(y_1 \mid x_1,x_2)
]

[
\log \pi_\theta(y_2 \mid x_1,x_2,y_1)
]

[
\log \pi_\theta(y_3 \mid x_1,x_2,y_1,y_2)
]

这和推理阶段逐步 decode 得到的 log p 是同一个东西。

需要注意的是，prompt token 通常不参与 RL loss，但它们必须作为上下文参与 attention。也就是说：

• attention mask 不能把 prompt 屏蔽掉；
• loss mask / response mask 只是在计算 loss 时忽略 prompt 部分。

如果训练端真的只输入 response，那么算出来的是：

[
\log \pi_\theta(y_t \mid y_{<t})
]

这当然无法和推理阶段的：

[
\log \pi_\theta(y_t \mid x, y_{<t})
]

对齐。这种流程就是错的。

---

为什么一次训练前向可以等价

decoder-only causal LM 的序列概率分解为：

[
\pi_\theta(y \mid x)

\prod_{t=1}^{T}
\pi_\theta(y_t \mid x, y_{<t})
]

推理阶段是一个因子一个因子地采样：

p(y1 | x)
p(y2 | x, y1)
p(y3 | x, y1, y2)
...

训练阶段用 teacher forcing，把完整序列：

[prompt, response]

一次送入模型。由于 causal mask 的存在，每个位置只能看见自己左边的 token，不能看见未来 token，所以它可以并行计算所有条件概率。

以 token 序列：

s = [x1, x2, y1, y2, y3]

为例，模型输出 logits：

z0, z1, z2, z3, z4

其中：

z1 -> predict y1
z2 -> predict y2
z3 -> predict y3

最后一个 logits z4 是用来预测 y3 后面的下一个 token 的，通常不参与当前 response 的 log p 计算。

一个简化的 PyTorch 对齐逻辑如下：

# prompt_ids: [m]
# response_ids: [n]
# ids: [1, m+n]
ids = torch.cat([prompt_ids, response_ids], dim=0).unsqueeze(0)

logits = model(ids).logits          # [1, m+n, vocab]

# logits[:, i] predicts ids[:, i+1]
logp_next = logits[:, :-1].log_softmax(dim=-1)
targets = ids[:, 1:]

token_logp = logp_next.gather(
    dim=-1,
    index=targets.unsqueeze(-1)
).squeeze(-1)

# response 的第一个 token 在 full sequence 中的位置是 m
# 它由 logits 的位置 m-1 预测
resp_logp = token_logp[:, m - 1 : m - 1 + len(response_ids)]

这个 resp_logp 就应该和推理引擎在 rollout 时记录的 output token logprobs 对齐。

如果你使用 HuggingFace CausalLM 的 labels，常见做法是：

labels = ids.clone()
labels[:, :prompt_len] = -100

因为 HF CausalLM 内部通常会做 shift：用 logits[:, :-1] 预测 labels[:, 1:]。所以把原始 labels 中 prompt 部分置为 -100 后，第一个 response token 仍然会由最后一个 prompt token 的 logits 来预测。

---

prefill + decode 和 full forward 的关系

prefill + decode 不是另一种概率模型，它只是带 KV cache 的增量计算。

推理：

prefill(prompt) 生成 prompt 的 KV cache
使用最后一个 prompt 位置的 logits 采样 y1
decode(y1) 更新 KV cache，得到 y2 的 logits
decode(y2) 更新 KV cache，得到 y3 的 logits
...

训练 full forward：

一次性输入 [prompt, y1, y2, y3, ...]
用 causal mask 同时算出所有位置的 logits

KV cache 的作用只是避免重复计算历史 token 的 key/value。理论上：

full forward 的 prefix hidden states

和：

prefill + decode 逐步得到的 hidden states

应该一致。

因此，只要以下条件一致，二者的 per-token log p 应该接近相等：

• 模型权重一致；
• tokenizer 和 token ids 一致；
• attention mask 一致；
• position ids 一致；
• RoPE scaling、YaRN、NTK 等位置编码配置一致；
• logits 处理方式一致；
• dtype 和 kernel 数值误差在可接受范围内。

实际工程中不一定 bitwise 相等，尤其是 bf16、FlashAttention、PagedAttention、tensor parallel、FP8 KV cache 等场景，但差异应该很小。若出现系统性大偏差，就说明存在训推不一致。

---

RL 崩溃为什么和 log p 不一致有关

以 PPO / GRPO 类算法为例，训练时通常需要 old logprob 和 new logprob：

[
r_t =
\exp(
\log \pi_{\theta_{\text{new}}}(y_t \mid x,y_{<t})

\log \pi_{\theta_{\text{old}}}(y_t \mid x,y_{<t})
)
]

其中：

• old_logprob：rollout 时生成该 token 的策略概率；
• new_logprob：训练时当前 actor 对同一个 token 的概率；
• (r_t)：importance ratio。

如果刚同步完 actor，且还没有做 optimizer step，那么理论上：

new_logprob ≈ old_logprob
r_t ≈ 1

如果此时就出现较大偏差，比如：

new_logprob - old_logprob

大面积偏离 0，那么 PPO ratio 会被错误放大或缩小，clip、KL、advantage 加权都会失真，训练就可能崩溃。

所以这里说的“训推不一致”，通常不是指“训练输入 response，推理输入 prompt”这种概念差异，而是指：

rollout/inference engine 记录的 log p

和：

training engine 对同一批 prompt+response 重算的 log p

不一致。

---

常见的不一致来源

1. token 序列不一致

这是最常见的问题。

需要确保训练端使用的不是重新 detokenize 再 retokenize 的文本，而是 rollout 时真实生成的 token ids：

prompt_token_ids + generated_response_token_ids

常见坑包括：

• BOS 是否自动添加；
• chat template 是否一致；
• 是否使用 add_generation_prompt=True；
• assistant header 是否被当成 prompt 还是 response；
• EOS token 是否包含在 response 中；
• stop string 和 stop token 的处理是否一致；
• 训练端是否重新拼文本后再 tokenize，导致边界 token 变化。

2. shift 和 mask 错位

response 第一个 token 的 logprob 来自最后一个 prompt token 的 logits。

如果 prompt 长度为 (m)，response 长度为 (n)，那么 response logprob 对应：

logits[m-1 : m-1+n]

而不是：

logits[m : m+n]

常见错误是 off-by-one。

3. logits processor 不一致

推理时可能使用：

• temperature；
• top-p；
• top-k；
• repetition penalty；
• min length；
• bad words；
• forced EOS；
• stop token suppression。

如果推理记录的是处理后的 logprob，而训练端用的是原始 logits 的 logprob，就会不一致。

例如 temperature 为 (\tau) 时，采样分布是：

[
\text{softmax}(z / \tau)
]

而不是：

[
\text{softmax}(z)
]

所以要么两边都用 raw logits，要么两边都应用相同的 temperature。调试时建议先关闭所有 processor：

temperature = 1
top_p = 1
top_k disabled
repetition_penalty = 1

4. padding、position id、packing 不一致

训练端常用 padding / sequence packing，推理端常用动态 batching / paged attention。需要确保：

• pad token 不参与 attention；
• prompt 不被 loss mask 误当成 attention mask 屏蔽；
• position ids 计算一致；
• packed sequence 之间不能互相 attention；
• RoPE scaling 配置一致；
• left padding / right padding 不导致 position id 差异。

5. 模型权重版本不一致

RL 系统通常有两个模型副本：

rollout engine actor
training engine actor

需要确认：

• 权重是否已经同步；
• LoRA adapter 是否一致；
• rollout engine 是否加载了 merge 后权重；
• tensor parallel 切分是否一致；
• 是否存在异步 rollout 的 stale policy；
• 量化权重、FP8 KV cache 是否引入较大误差。

6. 数值 kernel 差异

full forward 和 decode 可能使用不同 kernel：

• FlashAttention；
• PagedAttention；
• fused RMSNorm；
• fused softmax；
• bf16/fp16/fp32；
• tensor parallel all-reduce；
• vocab parallel softmax。

小误差正常，但大面积系统性误差不正常。

经验上，fp32/eager 模式可以非常接近；bf16、不同 attention kernel 下可能有 (10^{-3}) 到 (10^{-2}) 量级差异。这个不是硬标准，但如果偏差明显更大，或者 ratio 明显偏离 1，就要排查。

---

建议的排查流程

可以按下面顺序做最小化一致性测试：

1. 固定一批 prompt。
2. rollout engine 生成 response，并保存：
   • prompt_token_ids
   • response_token_ids
   • old_logprobs
   • attention_mask
   • position_ids
3. 关闭 temperature、top-p、top-k、repetition penalty 等 logits processor。
4. 确认训练端和推理端使用同一份权重。
5. 训练端用：

input_ids = prompt_token_ids + response_token_ids

重算 per-token logprob。

6. 对齐 response 部分：

train_logprob[i] 对齐 infer_logprob[i]

其中：

train_logprob[i]
=
log p(response_ids[i] | prompt_ids, response_ids[:i])

7. 比较：

delta = train_logprob - infer_logprob
ratio = torch.exp(delta)

期望：

delta ≈ 0
ratio ≈ 1

如果不满足，优先检查：

token ids -> chat template -> shift -> mask -> logits processor -> position ids -> 权重同步 -> dtype/kernel

一句话总结：

推理是逐步生成同一条 response，训练是对这条固定 response 做 teacher-forcing 评分；二者对齐的是每个 token 的条件 log probability，而不是表面上的输入 API。prefill+decode 与一次 full forward 在 causal mask 下理论等价，工程上的不一致才是需要排查的核心。